import cv2 as cv
import numpy as np


def sigmoid(x):
    return 1/(1+np.exp(-x))


def decode_box(x):
    """输出解码的numpy版本，用于onnx"""
    x = [x[..., :300], x[..., 300:1500], x[..., 1500:]]
    x = [x[0].reshape(-1, 5, 15, 20), x[1].reshape(-1, 5, 30, 40), x[2].reshape(-1, 5, 60, 80)]
    predictions = []
    for i in range(3):
        predict = x[i].transpose(0, 2, 3, 1)  # 批大小,高,宽,预测结果 N,H,W,P
        img_height, img_width = predict.shape[1], predict.shape[2]
        grid_x = np.expand_dims(np.arange(0, img_width), axis=0).repeat(img_height, 0)
        grid_y = np.expand_dims(np.arange(0, img_height), axis=0).repeat(img_width, 0).transpose(1, 0)
        predict[..., 0] = sigmoid(predict[..., 0]) + grid_x  # 横坐标
        predict[..., 1] = sigmoid(predict[..., 1]) + grid_y  # 纵坐标
        predict[..., 2] = 7 * np.exp(predict[..., 2])  # 宽度
        predict[..., 3] = 10 * np.exp(predict[..., 3])  # 高度
        predict[..., 4] = sigmoid(predict[..., 4])  # 置信度
        predictions.append(predict)
    return tuple(predictions)


def image_generate_loc(img, label):
    """在图像上显示结果
    :param img:网络输入的原图像
    :param label:网络经解码后的输出
    """
    stride = 32, 16, 8
    cnt = [0, 0, 0]
    predict = []
    for fm in range(3):
        predict.append(label[fm].reshape(-1, 5))
        cnt[fm] = predict[fm].shape[0]
    predict = np.vstack(predict)
    index = predict[:, 4].argmax()
    predict = predict[index]
    if predict[4] < 0.4:
        return img
    corner = predict[:4].copy()
    corner[0] = predict[0] - predict[2]/2
    corner[1] = predict[1] - predict[3]/2
    corner[2] = predict[0] + predict[2]/2
    corner[3] = predict[1] + predict[3]/2
    fm = 0
    while index > 0:
        index -= cnt[fm]
        fm += 1
    fm -= 1
    corner = np.round(corner*stride[fm]).astype('int32')
    img = cv.rectangle(img, (corner[0], corner[1]), (corner[2], corner[3]), color=(0, 0, 255), thickness=2)
    pred_conf = '{:.4f}'.format(predict[4])
    img = cv.putText(img, pred_conf, (corner[0], corner[1]), cv.FONT_HERSHEY_PLAIN, 1.2, (0, 0, 255), 1)
    return img


print('onnx detect.')
net = cv.dnn.readNetFromONNX('zhnnet.onnx')
image_origin = cv.imread('E:/dataset/instrument/pos1.png')
assert image_origin is not None, 'Image does not exit.'
image = image_origin.transpose(2, 0, 1) / 256
image = np.expand_dims(image, axis=0)
net.setInput(image)
pred = net.forward()
pred = decode_box(pred)
image_origin = image_generate_loc(image_origin, pred)
cv.imshow('test', image_origin)
cv.waitKey(0)
